﻿@using Blazor.Diagrams.Components
@using Blazor.Diagrams.Core;
@using Blazor.Diagrams.Core.Geometry;
@using Blazor.Diagrams.Core.Models;
@using Blazor.Diagrams;
@using Blazor.Diagrams.Core.Options;
@using Blazor.Diagrams.Core.PathGenerators;
@using Blazor.Diagrams.Core.Routers;
@using Blazor.Diagrams.Options;
@using HardHatCore.HardHatC2Client.Models
@using HardHatCore.HardHatC2Client.Pages
@using System.Security.Cryptography
@using HardHatCore.HardHatC2Client.Plugin_BaseClasses;
@using HardHatCore.HardHatC2Client.Plugin_Interfaces;

<MudButton style="z-index: 999;" @onclick="() =>ResetDiagram()">Reset Diagram</MudButton>
<MudButton style="z-index: 999;" @onclick="() => Diagram.ZoomToFit(50)">Zoom to fit</MudButton>

<CascadingValue Value="Diagram">
    <DiagramCanvas Class="grid">
            <Widgets>
                <NavigatorWidget Width="100" Height="100" DefaultStyle="true"></NavigatorWidget>
            </Widgets>
        </DiagramCanvas>
</CascadingValue>

@code {
    

    
    public class ImplantNode : NodeModel
    {
        public ImplantNode(Point position = null) : base(position) { }

        public ExtImplant_Base _implant { get; set; }
        public string Name { get; set;}
    }

    public class TeamServerNode : NodeModel
    {
        public TeamServerNode(Point position = null) : base(position) { }

        public string Name { get; set;}
    }

    private static BlazorDiagram Diagram { get; set; }
    private static Dictionary<string,Point> http_even_positions = new Dictionary<string, Point>();
    private static Dictionary<string,Point> http_odd_positions = new Dictionary<string, Point>();
    public static bool ShouldRenderDiagram = true;


    private ImplantNode NewImplantNode(double x, double y, ExtImplant_Base implant)
    {
        var node = new ImplantNode(new Point(x, y));
        node.Size = new Blazor.Diagrams.Core.Geometry.Size(200, 200);
        node._implant = implant;
        node.AddPort(PortAlignment.Left);
        node.AddPort(PortAlignment.Right);
        return node;
    }

    private TeamServerNode NewTSNode(double x, double y)
    {
        var node = new TeamServerNode(new Point(x, y));
        node.Size = new Blazor.Diagrams.Core.Geometry.Size(200, 200);
        node.AddPort(PortAlignment.Left);
        node.AddPort(PortAlignment.Right);
        return node;
    }
    
    public void ResetDiagram()
    {
        ShouldRenderDiagram = true;
        StateHasChanged();
    }

    protected override void OnInitialized()
    {
        try
        {
            base.OnInitialized();
                Diagram = new BlazorDiagram(new BlazorDiagramOptions
                {
                    GridSize = 50,
                });

    //Diagram.Options.DefaultNodeComponent = typeof(SvgNodeWidget);
                Diagram.RegisterComponent<ImplantNode, ImplantSvgNodeWidget>();
                Diagram.RegisterComponent<TeamServerNode, TeamServerSvgNodeWidget>();
    //foreach engineer in the list of engineers in the Engineers page add a node to the diagram
                int http_count = 0;

                TeamServerNode node = NewTSNode(0, 600);
                node.Name = "TeamServer";
                http_even_positions.Add("TeamServer", new Point(0, 600));
                http_odd_positions.Add("TeamServer", new Point(0, 600));
                Diagram.Nodes.Add(node);
                http_count++;

            foreach (ExtImplant_Base implant in Implants.ImplantDisplayList)
                {
    //if engineer.ConnectionTye is smb or tcp skip it because we want to map all the http/https connections first 
                    if (implant.ConnectionType.Equals("smb bind", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("smb reverse", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("tcp bind", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("tcp reverse", StringComparison.OrdinalIgnoreCase))
                    {
                        continue;
                    }

                    ImplantNode implant_node = NewImplantNode(0, 0, implant);
                    Point newPosition = new Point(0, 0);
    //taking the lastPostion if count is odd and lastPostion.X equals 400 add 75 to the x and -75 to the y, otherwise just add -75 to y,
    //if it is even and lastPostion.X is equal to 400 add 75 to the x and 75 to the y, if y would be greater then 800 or less then 0 reset it to 50 and add 75 to the x
                    if (http_count % 2 == 0)
                    {
                        Point lastPosition = http_even_positions.Last().Value;
                        if (lastPosition.Y == 600)
                        {
    //moving from the start over 200 so there is a 50px gap between the first node and the teamserver
                            newPosition = new Point(lastPosition.X + 250, lastPosition.Y + 200);
                        }
                        else
                        {
    //moving down 250 in the new row
                            newPosition = new Point(lastPosition.X, lastPosition.Y + 200);
                        }
                        if (newPosition.Y > 1050)
                        {
    //reset the chart height to 600, and move over 250 to the right to start a new row 
                            newPosition = new Point(lastPosition.X + 250, 600);
                            implant_node = NewImplantNode(newPosition.X, newPosition.Y, implant);
                        implant_node.Name = $"{implant.Metadata.ProcessId}@{implant.Metadata.Address}";
                        http_even_positions.Add(implant.Metadata.Id, newPosition);
                            Diagram.Nodes.Add(implant_node);
                            if (implant.ConnectionType.Equals("http", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("https", StringComparison.OrdinalIgnoreCase))
                            {
                                TeamServerNode node1 = (TeamServerNode)Diagram.Nodes.OfType<TeamServerNode>().FirstOrDefault(x => x.Name.Equals("TeamServer"));
                            ImplantNode node2 = (ImplantNode)Diagram.Nodes.OfType<ImplantNode>().FirstOrDefault(x => x.Name.Equals($"{implant.Metadata.ProcessId}@{implant.Metadata.Address}"));
                                Diagram.Links.Add(new LinkModel(node1, node2)
                                {
                                    SourceMarker = LinkMarker.Arrow,
                                    TargetMarker = LinkMarker.Square,
                                    Router = new NormalRouter(),
                                    PathGenerator = new StraightPathGenerator(),
                                    Color = "green"
                                });
                            }
                        }
                        else
                        {
                            implant_node = NewImplantNode(newPosition.X, newPosition.Y, implant);
                        implant_node.Name = $"{implant.Metadata.ProcessId}@{implant.Metadata.Address}";
                        http_even_positions.Add(implant.Metadata.Id, newPosition);
                            Diagram.Nodes.Add(implant_node);
                            if (implant.ConnectionType.Equals("http", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("https", StringComparison.OrdinalIgnoreCase))
                            {
                                TeamServerNode node1 = (TeamServerNode)Diagram.Nodes.OfType<TeamServerNode>().FirstOrDefault(x => x.Name.Equals("TeamServer"));
                            ImplantNode node2 = (ImplantNode)Diagram.Nodes.OfType<ImplantNode>().FirstOrDefault(x => x.Name.Equals($"{implant.Metadata.ProcessId}@{implant.Metadata.Address}"));
                                Diagram.Links.Add(new LinkModel(node1, node2)
                                {
                                    SourceMarker = LinkMarker.Arrow,
                                    TargetMarker = LinkMarker.Square,
                                    Router = new NormalRouter(),
                                    PathGenerator = new StraightPathGenerator(),
                                    Color = "green"

                                });
                            }
                        }
                    }
                    else
                    {
                        Point lastPosition = http_odd_positions.Last().Value;
                        if (lastPosition.Y == 600)
                        {
                            newPosition = new Point(lastPosition.X + 250, lastPosition.Y - 200);

                        }
                        else
                        {
    //moving up 150 in the new row
                            newPosition = new Point(lastPosition.X, lastPosition.Y - 200);
                        }
                        if (newPosition.Y < 0)
                        {
    //reset the chart height to 600
                            newPosition = new Point(lastPosition.X + 250, 600);
                            implant_node = NewImplantNode(newPosition.X, newPosition.Y, implant);
                        implant_node.Name = $"{implant.Metadata.ProcessId}@{implant.Metadata.Address}";
                        http_odd_positions.Add(implant.Metadata.Id, newPosition);
                            Diagram.Nodes.Add(implant_node);
                            if (implant.ConnectionType.Equals("http", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("https", StringComparison.OrdinalIgnoreCase))
                            {
                                TeamServerNode node1 = (TeamServerNode)Diagram.Nodes.OfType<TeamServerNode>().FirstOrDefault(x => x.Name.Equals("TeamServer"));
                            ImplantNode node2 = (ImplantNode)Diagram.Nodes.OfType<ImplantNode>().FirstOrDefault(x => x.Name.Equals($"{implant.Metadata.ProcessId}@{implant.Metadata.Address}"));
                                Diagram.Links.Add(new LinkModel(node1, node2)
                                {
                                    SourceMarker = LinkMarker.Arrow,
                                    TargetMarker = LinkMarker.Square,
                                    Router = new NormalRouter(),
                                    PathGenerator = new StraightPathGenerator(),
                                    Color = "green"

                                });
                            }
                        }
                        else
                        {
                            implant_node = NewImplantNode(newPosition.X, newPosition.Y, implant);
                        implant_node.Name = $"{implant.Metadata.ProcessId}@{implant.Metadata.Address}";
                        http_odd_positions.Add(implant.Metadata.Id, newPosition);
                            Diagram.Nodes.Add(implant_node);
                            if (implant.ConnectionType.Equals("http", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("https", StringComparison.OrdinalIgnoreCase))
                            {
                                TeamServerNode node1 = (TeamServerNode)Diagram.Nodes.OfType<TeamServerNode>().FirstOrDefault(x => x.Name.Equals("TeamServer"));
                            ImplantNode node2 = (ImplantNode)Diagram.Nodes.OfType<ImplantNode>().FirstOrDefault(x => x.Name.Equals($"{implant.Metadata.ProcessId}@{implant.Metadata.Address}"));
                                Diagram.Links.Add(new LinkModel(node1, node2)
                                {
                                    SourceMarker = LinkMarker.Arrow,
                                    TargetMarker = LinkMarker.Square,
                                    Router = new NormalRouter(),
                                    PathGenerator = new StraightPathGenerator(),
                                    Color = "green"

                                });
                            }
                        }
                    }
                    http_count++;
                }

    //for each Engineer in the list of Engineers, if it has the connection type of http or https, skip it, otherwise the engineer is a p2p connection and its external address value is the name of a node in the diagram, which will be used for the connection, 
    // the connection will be red for tcp and blue for smb, if it is a bind connection the sourceMarker will be a square and the targetMarker will be an arrow, if it is a reverse connection the sourceMarker will be an arrow and the targetMarker will be a square
            foreach (ExtImplant_Base implant in Implants.ImplantDisplayList)
                {
                    if (implant.ConnectionType.Equals("http", StringComparison.OrdinalIgnoreCase) || implant.ConnectionType.Equals("https", StringComparison.OrdinalIgnoreCase))
                    {
                        continue;
                    }

                    Point newPosition = new Point(0, 0);
    //take the engineers external address and use it to find the node in the diagram
                    ImplantNode node1 = (ImplantNode)Diagram.Nodes.OfType<ImplantNode>().FirstOrDefault(x => x.Name.Equals($"{implant.ExternalAddress}"));
    //update newPostion to be the position of the node with +250 to x and +0 to y, this will be the position of the new node, if any node is already at that position, move up 150 in the y axis
                    newPosition = new Point(node1.Position.X + 250, node1.Position.Y);

    //check that none of the Nodes in the diagram have the same position as the new position
                    foreach (NodeModel position_test_node in Diagram.Nodes)
                    {
                        if (position_test_node.Position.Equals(newPosition))
                        {
                            newPosition = new Point(node.Position.X, node.Position.Y - 150);
                        }
                    }
    //if the new position is less than 0, reset the chart height to 600 and test if it is overlapping again but move up 150 in the y axis
                    if (newPosition.Y < 0)
                    {
                        newPosition = new Point(node1.Position.X, 600);
                        foreach (NodeModel position_test_node in Diagram.Nodes)
                        {
                            if (position_test_node.Position.Equals(newPosition))
                            {
                                newPosition = new Point(node.Position.X, node.Position.Y + 150);
                            }
                        }
                    }

    //if the new postion Y is greater then 1050 then reset the chart height to 600, move to the right 250 in the x axis and move up 150 in the y axis
                    if (newPosition.Y > 1050)
                    {
                        newPosition = new Point(node1.Position.X + 250, 600);
                        foreach (NodeModel position_test_node in Diagram.Nodes)
                        {
                            if (position_test_node.Position.Equals(newPosition))
                            {
                                newPosition = new Point(node.Position.X, node.Position.Y + 150);
                            }
                        }
                    }

    //create a new node with the new position and the engineer object
                    ImplantNode node2 = NewImplantNode(newPosition.X, newPosition.Y, implant);
    //add the new node to the diagram
                    Diagram.Nodes.Add(node2);

                string eng_connection_type = implant.ConnectionType + " " + Implants.FindManager(implant, implant.Metadata.ManagerName);

                    if (eng_connection_type.Equals("tcp bind", StringComparison.OrdinalIgnoreCase))
                    {
    //add a link between the two nodes, the sourceMarker will be a square and the targetMarker will be an arrow
                        Diagram.Links.Add(new LinkModel(node1, node2)
                        {
                            SourceMarker = LinkMarker.Square,
                            TargetMarker = LinkMarker.Arrow,
                            Router = new NormalRouter(),
                            PathGenerator = new StraightPathGenerator(),
                            Color = "red"
                        });
                    }
                    if (eng_connection_type.Equals("tcp reverse", StringComparison.OrdinalIgnoreCase))
                    {
    //add a link between the two nodes, the sourceMarker will be a square and the targetMarker will be an arrow
                        Diagram.Links.Add(new LinkModel(node1, node2)
                        {
                            SourceMarker = LinkMarker.Arrow,
                            TargetMarker = LinkMarker.Square,
                            Router = new NormalRouter(),
                            PathGenerator = new StraightPathGenerator(),
                            Color = "red"
                        });

                    }
                    if (eng_connection_type.Equals("smb bind", StringComparison.OrdinalIgnoreCase))
                    {
    //add a link between the two nodes, the sourceMarker will be a square and the targetMarker will be an arrow
                        Diagram.Links.Add(new LinkModel(node1, node2)
                        {
                            SourceMarker = LinkMarker.Square,
                            TargetMarker = LinkMarker.Arrow,
                            Router = new NormalRouter(),
                            PathGenerator = new StraightPathGenerator(),
                            Color = "blue"
                        });

                    }
                    if (eng_connection_type.Equals("smb reverse", StringComparison.OrdinalIgnoreCase))
                    {
    //add a link between the two nodes, the sourceMarker will be a square and the targetMarker will be an arrow
                        Diagram.Links.Add(new LinkModel(node1, node2)
                        {
                            SourceMarker = LinkMarker.Arrow,
                            TargetMarker = LinkMarker.Square,
                            Router = new NormalRouter(),
                            PathGenerator = new StraightPathGenerator(),
                            Color = "blue"
                        });

                    }

                }
        }
        catch (Exception ex)
        {
            Console.WriteLine(ex.Message);
            Console.WriteLine(ex.StackTrace);
        }
    }

}
